package com.github.cyf.util.video;

import org.bytedeco.javacv.FFmpegFrameGrabber;
import org.bytedeco.javacv.Frame;
import org.bytedeco.javacv.Java2DFrameConverter;
import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;

/**
 * @author chenyifan
 * @create 2024-10-05 22:17
 */
public class VideoUtil {
    private static final String IMAGE_TYPE = ".jpg";

    /**
     * 获取视频帧
     * @param videoPath 视频路径
     * @param saveDir 保存目录
     * @return
     */
    public static List<File> getVideoFrames(String videoPath, String saveDir) {
        File video = new File(videoPath);
        List<File> images = new ArrayList<>();
        try (FFmpegFrameGrabber frameGrabber = new FFmpegFrameGrabber(video)) {
            frameGrabber.start();
            int size = frameGrabber.getLengthInFrames();
            for (int i = 0; i< size; i++) {
                // 获取帧
                Frame frame = frameGrabber.grabImage();
                String framePath = Paths.get(saveDir, String.format("%05d", i + 1) + IMAGE_TYPE).toString();
                File image = frameToImage(frame, framePath);
                if (image != null) {
                    images.add(image);
                }
            }
            return images;
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 帧转为图片
     * @param frame 帧
     * @return
     */
    private static File frameToImage(Frame frame, String filePath) {
        Java2DFrameConverter converter = new Java2DFrameConverter();
        BufferedImage image = converter.convert(frame);
        if (image == null) {
            return null;
        }
        File file = new File(filePath);
        if (!file.getParentFile().exists()) {
            file.getParentFile().mkdirs();
        }
        try {
            ImageIO.write(image, IMAGE_TYPE.substring(1), file);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        return file;
    }
}
